//==========================================================================
// Copyright (c) 2000-2008,  Elastos, Inc.  All Rights Reserved.
//==========================================================================

#ifdef _MSVC

#include <winhack.h>
#include <elasys_server.h>
#include <stdio.h>
#include <stdlib.h>
#include <cmalloc.h>
//#include <pal.h>

// WIN32 debug API
typedef struct _IMAGEHLP_LINE {
  DWORD SizeOfStruct;
  PVoid Key;
  DWORD LineNumber;
  PCHAR FileName;
  DWORD Address;
} IMAGEHLP_LINE, *PIMAGEHLP_LINE;

typedef BOOL (__stdcall *pSymInitialize)(
  HANDLE hProcess,
  PSTR UserSearchPath,
  BOOL fInvadeProcess
);

typedef BOOL (__stdcall *pSymGetLineFromAddr)(
  HANDLE hProcess,
  DWORD dwAddr,
  PDWORD pdwDisplacement,
  PIMAGEHLP_LINE Line
);

typedef struct _IMAGEHLP_SYMBOL {
  DWORD SizeOfStruct;
  DWORD Address;
  DWORD Size;
  DWORD Flags;
  DWORD MaxNameLength;
  CHAR  Name[1];
} IMAGEHLP_SYMBOL, *PIMAGEHLP_SYMBOL;

typedef BOOL (__stdcall *pSymGetSymFromAddr)(
  HANDLE hProcess,
  DWORD Address,
  PDWORD Displacement,
  PIMAGEHLP_SYMBOL Symbol
);

typedef enum {
    SymNone,
    SymCoff,
    SymCv,
    SymPdb,
    SymExport,
    SymDeferred,
    SymSym                  // .sym file
} SYM_TYPE;

typedef struct _IMAGEHLP_MODULE {
  DWORD    SizeOfStruct;
  DWORD    BaseOfImage;
  DWORD    ImageSize;
  DWORD    TimeDateStamp;
  DWORD    CheckSum;
  DWORD    NumSyms;
  SYM_TYPE SymType;
  CHAR    ModuleName[32];
  CHAR    ImageName[256];
  CHAR    LoadedImageName[256];
} IMAGEHLP_MODULE, *PIMAGEHLP_MODULE;

typedef  BOOL (__stdcall *pSymGetModuleInfo)(
  HANDLE hProcess,
  DWORD dwAddr,
  PIMAGEHLP_MODULE ModuleInfo
);

typedef DWORD (__stdcall *pSymSetOptions)(
  DWORD SymOptions
);

typedef BOOL (__stdcall *pSymCleanup)(
  HANDLE hProcess
);

static HANDLE s_hProcess;
BOOL g_bSymbolInit = FALSE;

pSymInitialize g_pSymInitialize;
pSymGetLineFromAddr g_pSymGetLineFromAddr;
pSymGetSymFromAddr g_pSymGetSymFromAddr;
pSymGetModuleInfo g_pSymGetModuleInfo;
pSymSetOptions g_pSymSetOptions;
pSymCleanup g_pSymCleanup;

EXTERN_C wchar_t * CDECL _GetEnvironment(const wchar_t *envstring);
void GetStackFuncPath(wchar_t *wscPathName)
{
    wchar_t *env = _GetEnvironment(L"ELASTOS_ROOT");
    if (NULL != env) {
        wcscpy(wscPathName, env);
        wcscat(wscPathName, L"\\build\\tools\\dbghelp.dll");
    }
    else {
        wcscpy(wscPathName, L"dbghelp.dll");
    }
}
EXTERN_C BOOL InitStackFunc()
{
    if (NULL != g_pSymInitialize)
        return TRUE;
    HMODULE hInst;
    BOOL ret = TRUE;
    s_hProcess = GetCurrentProcess();
    if (INVALID_HANDLE_VALUE == s_hProcess) {
        ret  = FALSE;
    }
    wchar_t wscPathName[MAX_PATH];
    GetStackFuncPath(wscPathName);
    hInst = LoadLibraryW(wscPathName);
    if (INVALID_HANDLE_VALUE == hInst) {
        return FALSE;
    }
    g_pSymInitialize = (pSymInitialize)GetProcAddress(hInst,"SymInitialize");
    g_pSymGetLineFromAddr = (pSymGetLineFromAddr)GetProcAddress(hInst,"SymGetLineFromAddr");
    g_pSymGetSymFromAddr = (pSymGetSymFromAddr)GetProcAddress(hInst,"SymGetSymFromAddr");
    g_pSymGetModuleInfo = (pSymGetModuleInfo)GetProcAddress(hInst,"SymGetModuleInfo");
    g_pSymSetOptions = (pSymSetOptions)GetProcAddress(hInst,"SymSetOptions");
    g_pSymCleanup = (pSymCleanup)GetProcAddress(hInst,"SymCleanup");
    if (NULL == g_pSymInitialize ||
        NULL == g_pSymGetLineFromAddr ||
        NULL == g_pSymGetSymFromAddr ||
        NULL == g_pSymGetModuleInfo ||
        NULL == g_pSymSetOptions ||
        NULL == g_pSymCleanup) {
        ret = FALSE;
    }
    return ret;
}

// BUGBUG: It is only fit msvc debug mode
EXTERN_C void GetCallerStack(Address *pCallstk, int level, Address limit)
{
    int i;
    ULONG chksum = 0;
    Address *pStkFrame;//0:esp 1:eip
    __asm mov pStkFrame, ebp;
    pStkFrame = (Address*)pStkFrame[0]; // Jump over the "CoTaskMemAlloc" Call
    Address pStkBase = (Address)pStkFrame;

    //try {
    for (i = 0; i < level; i++) {
        if ((Address)pStkFrame < pStkBase
            || (Address)pStkFrame >= limit)
            break;
//            printf("[ebp] %08x [eip1] %08x [eip2] %08x [eip3] %08x [eip4] %08x\n", pStkFrame[0], pStkFrame[1], pStkFrame[2], pStkFrame[3], pStkFrame[4]);
        pCallstk[i] = pStkFrame[1];
        chksum += pCallstk[i];
        pStkFrame = (Address*)pStkFrame[0];
    }
    //}
    //catch (...) {
    //
    //}
    if (i < level) {
        Address StackBottom = pCallstk[i];
        if (StackBottom < 0x40000 || StackBottom > 0xF0000000) {
            // Set invalid address to be NULL
            // BUGBUG: Are these conditions enough?
            pCallstk[i] = NULL;
        }
    }
    pCallstk[level] = chksum;
}
EXTERN_C BOOL InitSymbols();

EXTERN_C BOOL DumpStack(Address *pCallstk, int level)
{
    InitSymbols();

    Address CallAddr;
    for (int i = 0; i < level; i++) {
        CallAddr = pCallstk[i];

        if (!CallAddr) break;

        IMAGEHLP_LINE ImgLine;
        ImgLine.SizeOfStruct = sizeof(IMAGEHLP_LINE);
        ImgLine.Address = (DWORD)CallAddr;
        DWORD dwDisplacement = 0;
        int ret = g_pSymGetLineFromAddr(s_hProcess, CallAddr, &dwDisplacement, &ImgLine);

        if (ret) {
            wprintf(L"   0x%08x ", CallAddr);
            DWORD Displacement;
            BYTE SymbolBuffer[sizeof(IMAGEHLP_SYMBOL) + 512];
            PIMAGEHLP_SYMBOL pSymbol = (PIMAGEHLP_SYMBOL)SymbolBuffer;
            pSymbol->SizeOfStruct = sizeof(SymbolBuffer);
            pSymbol->MaxNameLength = 512;
            ret = g_pSymGetSymFromAddr(s_hProcess, CallAddr, &Displacement, pSymbol);
            if (ret) {
                IMAGEHLP_MODULE module;
                module.SizeOfStruct = sizeof(IMAGEHLP_MODULE);
                g_pSymGetModuleInfo(s_hProcess, CallAddr, &module);
                wprintf(L"%S!%S+0x%03x  ", module.ModuleName, pSymbol->Name, Displacement);
            }
            wprintf(L" [ %S/%d ]\n", ImgLine.FileName, ImgLine.LineNumber);
        }
        else {
            wprintf(L"   0x%08x\n", CallAddr);
//            wprintf(L"   ERROR:0x%08x\n", GetLastError());
        }
    }
    wprintf(L"\n");

    return TRUE;
}

EXTERN_C BOOL GetLocationFromSym(DWORD CallAddr, wchar_t *szLoc)
{
    if (NULL == s_hProcess) {
        s_hProcess = GetCurrentProcess();
    }
    if (g_bSymbolInit) {
        g_pSymCleanup(s_hProcess);
    }
    g_bSymbolInit = FALSE;
    InitSymbols();
    PIMAGEHLP_SYMBOL pSymbol;
    IMAGEHLP_MODULE module;
    DWORD Displacement = 0;
    BOOL ret;
    BYTE SymbolBuffer[sizeof(IMAGEHLP_SYMBOL) + 256];

    pSymbol = (PIMAGEHLP_SYMBOL)SymbolBuffer;
    pSymbol->SizeOfStruct = sizeof(SymbolBuffer);
    pSymbol->MaxNameLength = 256;
    ret = g_pSymGetSymFromAddr(s_hProcess, CallAddr, &Displacement, pSymbol);
    if (!ret) goto ErrorExit;

    module.SizeOfStruct = sizeof(IMAGEHLP_MODULE);
    ret = g_pSymGetModuleInfo(s_hProcess, CallAddr, &module);
    if (!ret) goto ErrorExit;
    swprintf(szLoc, L"%S!%S", module.ModuleName, pSymbol->Name);
    return TRUE;

ErrorExit:
    szLoc[0] = L'\0';
    return FALSE;
}

EXTERN_C void CleanupSymbols()
{
    g_pSymCleanup(s_hProcess);
//    wprintf(L"CleanUpSymbols OK!\n");
}

#define SYMOPT_UNDNAME                  0x00000002
#define SYMOPT_LOAD_LINES               0x00000010

EXTERN_C BOOL InitSymbols()
{
    if (g_bSymbolInit) return TRUE;
    if (!InitStackFunc()) {
        printf("Can not init stack func succeed! Use stack command may cause mistake\n");
    }
    BOOL ret;
//    wprintf(L"DumpSymbols...\n");
    s_hProcess = GetCurrentProcess();
    wchar_t *wszT = _GetEnvironment(L"path");
    char SearchPath[_MAX_PATH];
    wcstombs(SearchPath, wszT, _MAX_PATH);
    ret = g_pSymInitialize(s_hProcess, SearchPath, TRUE);
    if (!ret) {
        wprintf(L"SymInitialize error\n");
        return FALSE;
    }
    ret = g_pSymSetOptions(SYMOPT_LOAD_LINES | SYMOPT_UNDNAME);
    if (!ret) {
        wprintf(L"SymSetOptions error\n");
        return FALSE;
    }

//    wprintf(L"DumpSymbols ok\n");

    g_bSymbolInit = TRUE;

    static BOOL s_bCleanup = FALSE;
    if (!s_bCleanup) {
        atexit(CleanupSymbols);
        s_bCleanup = TRUE;
    }

    return TRUE;
}
EXTERN_C TIME _GetTickCount();
/////////////////////////////////////////////////////////////////////
EXTERN_C void GetObjCallerStack(Address *pCallstk, Address *pCallobj, int level)
{
    int i;
    ULONG chksum = 0;
    Address *pStkFrame;//0:esp 1:eip
    __asm mov pStkFrame, ebp;
    pStkFrame = (Address*)pStkFrame[0];
    //try {
    pCallstk[0] = 0;
    pStkFrame = (Address*)pStkFrame[0];
    for (i = 1; i < level; i++) {
        pCallstk[i] = pStkFrame[1];
        chksum += pCallstk[i];
        pCallobj[i] = pStkFrame[2];
        pStkFrame = (Address*)pStkFrame[0];
    }
    //}
    //catch (...) {
    //
    //}
    if (i < level) {
        Address StackBottom = pCallstk[i];
        if (StackBottom < 0x40000 || StackBottom > 0xF0000000) {
            // Set invalid address to be NULL
            // BUGBUG: Are these conditions enough?
            pCallstk[i] = NULL;
            pCallobj[i] = NULL;
        }
    }

    pCallstk[level++] = chksum;
    *(TIME *)&pCallstk[level] = _GetTickCount();//stamp
}

EXTERN_C BOOL DumpObjStack(Address *pCallstk, Address *pCallobj, int level)
{
    InitSymbols();

    Address CallAddr;
    for (int i = 0; i < level; i++) {
        CallAddr = pCallstk[i];

        if (!CallAddr) break;

        IMAGEHLP_LINE ImgLine;
        ImgLine.SizeOfStruct = sizeof(IMAGEHLP_LINE);
        ImgLine.Address = (DWORD)CallAddr;
        DWORD dwDisplacement = 0;
        int ret = g_pSymGetLineFromAddr(s_hProcess, CallAddr, &dwDisplacement, &ImgLine);

        if (ret) {
            wprintf(L"   0x%08x : %08x ", CallAddr, pCallobj[i]);
            DWORD Displacement;
            BYTE SymbolBuffer[sizeof(IMAGEHLP_SYMBOL) + 512];
            PIMAGEHLP_SYMBOL pSymbol = (PIMAGEHLP_SYMBOL)SymbolBuffer;
            pSymbol->SizeOfStruct = sizeof(SymbolBuffer);
            pSymbol->MaxNameLength = 512;
            ret = g_pSymGetSymFromAddr(s_hProcess, CallAddr, &Displacement, pSymbol);
            if (ret) {
                IMAGEHLP_MODULE module;
                module.SizeOfStruct = sizeof(IMAGEHLP_MODULE);
                g_pSymGetModuleInfo(s_hProcess, CallAddr, &module);
                wprintf(L"%S!%S+0x%03x  ", module.ModuleName, pSymbol->Name, Displacement);
            }
            wprintf(L" [ %S/%d ]\n", ImgLine.FileName, ImgLine.LineNumber);
        }
        else {
            wprintf(L"   0x%08x\n", CallAddr);
//            wprintf(L"   ERROR:0x%08x\n", GetLastError());
        }
    }
    wprintf(L"\n");

    return TRUE;
}

#define MAX_REF 100
#define MAX_STACK_LEVEL 20

typedef struct {
    Address oldRef;
    ulong_t oldCnt;
    ulong_t refCnt;
    Address RefStkFrm[MAX_REF][MAX_STACK_LEVEL+5];
    Address RefObjFrm[MAX_REF][MAX_STACK_LEVEL+1];
}s_DebugRefInfo;

static s_DebugRefInfo ref[3];
static Address s_DebugObj = 0;
static unsigned long s_DebugSum = -1;

void DumpObjRefStack()
{
    if (s_DebugObj == 0) return ;

    ulong_t i = 0;
    wprintf(L"\n\n==== Enter Dump Obj %08x RefStack %d : %d ====\n",
            s_DebugObj, ref[1].refCnt, ref[2].refCnt);

    for (i = 0; i < ref[1].refCnt; i++) {
        wprintf(L"\n==== AddRef Stack %d, stamp=%I64x, chksum %08x ====\n",
                i, *(TIME *)&ref[1].RefStkFrm[i][MAX_STACK_LEVEL+1],
                ref[1].RefStkFrm[i][MAX_STACK_LEVEL]);
        DumpObjStack(ref[1].RefStkFrm[i], ref[1].RefObjFrm[i], MAX_STACK_LEVEL);
    }

    for (i = 0; i < ref[2].refCnt; i++) {
        wprintf(L"\n==== Release Stack %d, stamp=%I64x, chksum %08x ====\n",
                i, *(TIME *)&ref[1].RefStkFrm[i][MAX_STACK_LEVEL+1],
                ref[1].RefStkFrm[i][MAX_STACK_LEVEL]);
        DumpObjStack(ref[2].RefStkFrm[i], ref[1].RefObjFrm[i], MAX_STACK_LEVEL);
    }
    wprintf(L"\n==== Exit Dump Obj RefStack ====\n\n");
}

EXTERN_C Address __stdcall _RefHook(Address pObj, int n)
{
    unsigned long i = pObj - s_DebugObj;
    if (i > 8) goto exit;

    wprintf(L"\n=== RefHook: Object %x, Interface %x (%d) ===\n", s_DebugObj, pObj, n);

    if (i != 0) wprintf(L"MATCH %d!!!!\n", i);

    //Todo something
    if (ref[n].refCnt >= MAX_REF) goto exit;

    GetObjCallerStack(ref[n].RefStkFrm[ref[n].refCnt],
                ref[n].RefObjFrm[ref[n].refCnt], MAX_STACK_LEVEL);

    //Check Sum
    //__asm int 3
    if (s_DebugSum == (s_DebugSum & ref[n].RefStkFrm[ref[n].refCnt][MAX_STACK_LEVEL])) {
        DebugBreak();
    }

    //Fixup
    ref[n].RefStkFrm[ref[n].refCnt][0] = ref[n].oldRef;

    ref[n].refCnt++;

exit:
    return ref[n].oldRef;
}

#define REFHOOK(_n_) \
    __declspec( naked ) void RefHook##_n_() {\
        __asm push  ebp                 \
        __asm mov   ebp,esp             \
        __asm push  _n_                 \
        __asm push  [ebp+8] /*pIObject*/\
        __asm call  _RefHook            \
        __asm pop   ebp                 \
        __asm jmp   eax /*Back*/        \
    }

REFHOOK(1) REFHOOK(2)
    void SetRefCheck(PInterface pObj, unsigned long sum)
{
    DWORD dwOldProtect;
    if (s_DebugObj) {//Restore old obj
        Address *v =  (Address *)(*(Address *)s_DebugObj);
        v[1] = ref[1].oldRef;
        v[2] = ref[2].oldRef;
        s_DebugObj = 0;
    }

    Address *v =  (Address *)(*(Address *)pObj);
    printf("=== Trace IInterface %08x, Chksum %08x, AddRef %08x, Release %08x ===\n",
        pObj, sum, v[1], v[2]);

    VirtualProtect((void*)v, 12, PAGE_EXECUTE_READWRITE, &dwOldProtect);

    //Save & Hook vtable
    memset(ref, 0, sizeof(ref));
    s_DebugObj = (Address)pObj;
    s_DebugSum = sum;
    ref[1].oldRef = v[1];
    ref[2].oldRef = v[2];

    v[1] = (Address)RefHook1;//AddRef
    v[2] = (Address)RefHook2;//Release

    //Dummy First Ref
    _RefHook(s_DebugObj, 1);
}

#endif
